"""
结果聚合器 - 负责更新状态和发布事件

这个模块实现了结果聚合的核心功能：
- 监听GPU服务器完成事件
- 更新任务追踪记录
- 发布任务完成事件
- 提供统一的结果查询接口
"""

import asyncio
import logging
import threading
import time
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional, Union
from enum import Enum

from event_bus import EventBus, EventType, EventHandler, Event, get_event_bus

logger = logging.getLogger(__name__)


class TaskStatus(Enum):
    """任务状态枚举"""
    PENDING = "pending"
    ASSIGNED = "assigned"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"
    TIMEOUT = "timeout"


@dataclass
class TaskResult:
    """任务结果数据结构"""
    task_id: str
    status: TaskStatus
    result_url: Optional[str] = None
    error_message: Optional[str] = None
    processing_duration: Optional[float] = None
    server_id: Optional[str] = None
    completed_at: Optional[datetime] = None
    metadata: Optional[Dict[str, Any]] = None


@dataclass
class TaskRecord:
    """任务记录数据结构"""
    task_id: str
    client_ip: str
    submit_time: datetime
    assigned_server: Optional[str] = None
    server_endpoint: Optional[str] = None
    status: TaskStatus = TaskStatus.PENDING
    result_url: Optional[str] = None
    processing_duration: Optional[float] = None
    error_message: Optional[str] = None
    correlation_id: Optional[str] = None
    metadata: Optional[Dict[str, Any]] = None


class StateStorage(ABC):
    """状态存储抽象接口"""
    
    @abstractmethod
    async def get_task(self, task_id: str) -> Optional[TaskRecord]:
        """获取任务记录"""
        pass
    
    @abstractmethod
    async def update_task(self, task_record: TaskRecord) -> None:
        """更新任务记录"""
        pass
    
    @abstractmethod
    async def list_tasks(self, 
                        status: Optional[TaskStatus] = None,
                        server_id: Optional[str] = None,
                        limit: int = 100) -> List[TaskRecord]:
        """列出任务记录"""
        pass


class MemoryStateStorage(StateStorage):
    """内存状态存储实现"""
    
    def __init__(self):
        self.tasks: Dict[str, TaskRecord] = {}
        self.lock = threading.RLock()
        logger.info("内存状态存储已初始化")
    
    async def get_task(self, task_id: str) -> Optional[TaskRecord]:
        """获取任务记录"""
        with self.lock:
            return self.tasks.get(task_id)
    
    async def update_task(self, task_record: TaskRecord) -> None:
        """更新任务记录"""
        with self.lock:
            self.tasks[task_record.task_id] = task_record
            logger.debug(f"已更新任务记录: {task_record.task_id}")
    
    async def list_tasks(self, 
                        status: Optional[TaskStatus] = None,
                        server_id: Optional[str] = None,
                        limit: int = 100) -> List[TaskRecord]:
        """列出任务记录"""
        with self.lock:
            filtered_tasks = []
            
            for task in self.tasks.values():
                if status and task.status != status:
                    continue
                
                if server_id and task.assigned_server != server_id:
                    continue
                
                filtered_tasks.append(task)
                
                if len(filtered_tasks) >= limit:
                    break
            
            return filtered_tasks


class ResultAggregator(EventHandler):
    """结果聚合器"""
    
    def __init__(self, 
                 state_storage: Optional[StateStorage] = None,
                 event_bus: Optional[EventBus] = None):
        self.state_storage = state_storage or MemoryStateStorage()
        self.event_bus = event_bus or get_event_bus()
        self.lock = threading.RLock()
        
        # 订阅相关事件
        self.event_bus.subscribe(self, priority=10)
        
        logger.info("结果聚合器已初始化")
    
    def can_handle(self, event_type: EventType) -> bool:
        """判断是否能处理指定类型的事件"""
        return event_type in {
            EventType.TASK_COMPLETED,
            EventType.TASK_FAILED,
            EventType.TASK_CANCELLED,
            EventType.TASK_TIMEOUT
        }
    
    async def handle(self, event: Event) -> None:
        """处理事件"""
        try:
            if event.type == EventType.TASK_COMPLETED:
                await self._handle_task_completed(event)
            elif event.type == EventType.TASK_FAILED:
                await self._handle_task_failed(event)
            elif event.type == EventType.TASK_CANCELLED:
                await self._handle_task_cancelled(event)
            elif event.type == EventType.TASK_TIMEOUT:
                await self._handle_task_timeout(event)
        except Exception as e:
            logger.error(f"处理事件 {event.type.value} 时出错: {e}", exc_info=True)
    
    async def _handle_task_completed(self, event: Event) -> None:
        """处理任务完成事件"""
        task_id = event.data.get("task_id")
        if not task_id:
            logger.warning("任务完成事件缺少task_id")
            return
        
        # 获取任务记录
        task_record = await self.state_storage.get_task(task_id)
        if not task_record:
            logger.warning(f"未找到任务记录: {task_id}")
            return
        
        # 更新任务状态
        task_record.status = TaskStatus.COMPLETED
        task_record.result_url = event.data.get("result_url")
        task_record.processing_duration = event.data.get("processing_duration")
        task_record.completed_at = datetime.now(timezone.utc)
        task_record.metadata = event.data.get("metadata", {})
        
        # 保存更新
        await self.state_storage.update_task(task_record)
        
        # 发布任务完成事件
        await self.event_bus.publish(
            EventType.TASK_COMPLETED,
            source="result_aggregator",
            data={
                "task_id": task_id,
                "result_url": task_record.result_url,
                "processing_duration": task_record.processing_duration,
                "server_id": task_record.assigned_server,
                "metadata": task_record.metadata
            },
            correlation_id=task_record.correlation_id
        )
        
        logger.info(f"任务 {task_id} 已完成，结果URL: {task_record.result_url}")
    
    async def _handle_task_failed(self, event: Event) -> None:
        """处理任务失败事件"""
        task_id = event.data.get("task_id")
        if not task_id:
            logger.warning("任务失败事件缺少task_id")
            return
        
        # 获取任务记录
        task_record = await self.state_storage.get_task(task_id)
        if not task_record:
            logger.warning(f"未找到任务记录: {task_id}")
            return
        
        # 更新任务状态
        task_record.status = TaskStatus.FAILED
        task_record.error_message = event.data.get("error_message")
        task_record.processing_duration = event.data.get("processing_duration")
        task_record.completed_at = datetime.now(timezone.utc)
        task_record.metadata = event.data.get("metadata", {})
        
        # 保存更新
        await self.state_storage.update_task(task_record)
        
        # 发布任务失败事件
        await self.event_bus.publish(
            EventType.TASK_FAILED,
            source="result_aggregator",
            data={
                "task_id": task_id,
                "error_message": task_record.error_message,
                "processing_duration": task_record.processing_duration,
                "server_id": task_record.assigned_server,
                "metadata": task_record.metadata
            },
            correlation_id=task_record.correlation_id
        )
        
        logger.warning(f"任务 {task_id} 失败: {task_record.error_message}")
    
    async def _handle_task_cancelled(self, event: Event) -> None:
        """处理任务取消事件"""
        task_id = event.data.get("task_id")
        if not task_id:
            logger.warning("任务取消事件缺少task_id")
            return
        
        # 获取任务记录
        task_record = await self.state_storage.get_task(task_id)
        if not task_record:
            logger.warning(f"未找到任务记录: {task_id}")
            return
        
        # 更新任务状态
        task_record.status = TaskStatus.CANCELLED
        task_record.error_message = event.data.get("reason", "任务被取消")
        task_record.completed_at = datetime.now(timezone.utc)
        task_record.metadata = event.data.get("metadata", {})
        
        # 保存更新
        await self.state_storage.update_task(task_record)
        
        logger.info(f"任务 {task_id} 已取消")
    
    async def _handle_task_timeout(self, event: Event) -> None:
        """处理任务超时事件"""
        task_id = event.data.get("task_id")
        if not task_id:
            logger.warning("任务超时事件缺少task_id")
            return
        
        # 获取任务记录
        task_record = await self.state_storage.get_task(task_id)
        if not task_record:
            logger.warning(f"未找到任务记录: {task_id}")
            return
        
        # 更新任务状态
        task_record.status = TaskStatus.TIMEOUT
        task_record.error_message = "任务处理超时"
        task_record.completed_at = datetime.now(timezone.utc)
        task_record.metadata = event.data.get("metadata", {})
        
        # 保存更新
        await self.state_storage.update_task(task_record)
        
        logger.warning(f"任务 {task_id} 处理超时")
    
    async def create_task_record(self, 
                                task_id: str,
                                client_ip: str,
                                correlation_id: Optional[str] = None,
                                metadata: Optional[Dict[str, Any]] = None) -> TaskRecord:
        """创建任务记录"""
        task_record = TaskRecord(
            task_id=task_id,
            client_ip=client_ip,
            submit_time=datetime.now(timezone.utc),
            correlation_id=correlation_id,
            metadata=metadata or {}
        )
        
        await self.state_storage.update_task(task_record)
        
        # 发布任务创建事件
        await self.event_bus.publish(
            EventType.TASK_CREATED,
            source="result_aggregator",
            data={
                "task_id": task_id,
                "client_ip": client_ip,
                "metadata": metadata
            },
            correlation_id=correlation_id
        )
        
        logger.info(f"已创建任务记录: {task_id}")
        return task_record
    
    async def assign_task(self, 
                         task_id: str,
                         server_id: str,
                         server_endpoint: str) -> None:
        """分配任务到服务器"""
        task_record = await self.state_storage.get_task(task_id)
        if not task_record:
            raise ValueError(f"未找到任务记录: {task_id}")
        
        # 更新任务记录
        task_record.assigned_server = server_id
        task_record.server_endpoint = server_endpoint
        task_record.status = TaskStatus.ASSIGNED
        
        await self.state_storage.update_task(task_record)
        
        # 发布任务分配事件
        await self.event_bus.publish(
            EventType.TASK_ASSIGNED,
            source="result_aggregator",
            data={
                "task_id": task_id,
                "server_id": server_id,
                "server_endpoint": server_endpoint
            },
            correlation_id=task_record.correlation_id
        )
        
        logger.info(f"任务 {task_id} 已分配给服务器 {server_id}")
    
    async def start_task_processing(self, task_id: str) -> None:
        """开始任务处理"""
        task_record = await self.state_storage.get_task(task_id)
        if not task_record:
            raise ValueError(f"未找到任务记录: {task_id}")
        
        # 更新任务状态
        task_record.status = TaskStatus.PROCESSING
        
        await self.state_storage.update_task(task_record)
        
        # 发布任务开始事件
        await self.event_bus.publish(
            EventType.TASK_STARTED,
            source="result_aggregator",
            data={
                "task_id": task_id,
                "server_id": task_record.assigned_server
            },
            correlation_id=task_record.correlation_id
        )
        
        logger.info(f"任务 {task_id} 开始处理")
    
    async def get_task_result(self, task_id: str) -> Optional[TaskResult]:
        """获取任务结果"""
        task_record = await self.state_storage.get_task(task_id)
        if not task_record:
            return None
        
        return TaskResult(
            task_id=task_record.task_id,
            status=task_record.status,
            result_url=task_record.result_url,
            error_message=task_record.error_message,
            processing_duration=task_record.processing_duration,
            server_id=task_record.assigned_server,
            completed_at=task_record.completed_at,
            metadata=task_record.metadata
        )
    
    async def list_tasks(self, 
                        status: Optional[TaskStatus] = None,
                        server_id: Optional[str] = None,
                        limit: int = 100) -> List[TaskRecord]:
        """列出任务记录"""
        return await self.state_storage.list_tasks(status, server_id, limit)
    
    async def get_statistics(self) -> Dict[str, Any]:
        """获取统计信息"""
        all_tasks = await self.state_storage.list_tasks(limit=10000)
        
        stats = {
            "total_tasks": len(all_tasks),
            "status_counts": {},
            "server_counts": {},
            "average_processing_duration": 0.0,
            "success_rate": 0.0
        }
        
        total_duration = 0.0
        completed_count = 0
        successful_count = 0
        
        for task in all_tasks:
            # 状态统计
            status_key = task.status.value
            stats["status_counts"][status_key] = stats["status_counts"].get(status_key, 0) + 1
            
            # 服务器统计
            if task.assigned_server:
                server_key = task.assigned_server
                stats["server_counts"][server_key] = stats["server_counts"].get(server_key, 0) + 1
            
            # 处理时长统计
            if task.processing_duration:
                total_duration += task.processing_duration
                completed_count += 1
            
            # 成功率统计
            if task.status == TaskStatus.COMPLETED:
                successful_count += 1
        
        if completed_count > 0:
            stats["average_processing_duration"] = total_duration / completed_count
        
        if len(all_tasks) > 0:
            stats["success_rate"] = successful_count / len(all_tasks)
        
        return stats
